#include "SORTTracker.hxx"
#include "KalmanTracker.hxx"
#include "Hungarian.hxx"
#define DObj2TBox(obj) \
   Rect_<float>(obj.x, obj.y, obj.width, obj.height)

// Computes IOU between two bounding boxes
double GetIOU(Rect_<float> bb_test, Rect_<float> bb_gt)
{
	float in = (bb_test & bb_gt).area();
	float un = bb_test.area() + bb_gt.area() - in;

	if (un < DBL_EPSILON)
		return 0;

	return (double)(in / un);
}

SORTTracker::SORTTracker()
{
	_maxAge = 1;
	_minHits = 3;
	_iouThreshold = 0.3;
	KalmanTracker::kf_count = 0; // tracking id relies on this, so we have to reset it in each seq.

	_trkNum = 0;
	_detNum = 0;
}

//TrackingBox SORTTracker::detectedObject2TrackingBox(DetectedObject& dobj)
//{
//	TrackingBox tb;
//	tb.box.x = dobj.x;
//	tb.box.y = dobj.y;
//	tb.box.width = dobj.width;
//	tb.box.height = dobj.height;
//	tb.id = 0;
//	tb.frame = 0;
//
//	return tb;
//}

void SORTTracker::initialize(std::vector<DetectedObject>& retObjs)
{
	_trackers.clear();
	// initialize kalman trackers using first detections.
	for (unsigned int i = 0; i < retObjs.size(); i++)
	{
		DetectedObject dobj = retObjs.at(i);
		KalmanTracker trk = KalmanTracker(Rect_<float>(dobj.x,dobj.y,dobj.width,dobj.height));
		_trackers.push_back(trk);
	}
}

void SORTTracker::predict()
{
	// 3.1. get predicted locations from existing trackers.
	_predictedBoxes.clear();

	for (auto it = _trackers.begin(); it != _trackers.end();)
	{
		Rect_<float> pBox = (*it).predict();
		if (pBox.x >= 0 && pBox.y >= 0)
		{
			_predictedBoxes.push_back(pBox);
			it++;
		}
		else
		{
			it = _trackers.erase(it);
			//cerr << "Box invalid at frame: " << frame_count << endl;
		}
	}
}

void SORTTracker::associate(std::vector<DetectedObject>& retObjs)
{
	// 3.2. associate detections to tracked object (both represented as bounding boxes)
	// dets : detFrameData[fi]
	_trkNum = _predictedBoxes.size();
	_detNum = retObjs.size();

	_iouMatrix.clear();
	_iouMatrix.resize(_trkNum, vector<double>(_detNum, 0));

	for (unsigned int i = 0; i < _trkNum; i++) // compute iou matrix as a distance matrix
	{
		for (unsigned int j = 0; j < _detNum; j++)
		{
			// use 1-iou because the hungarian algorithm computes a minimum-cost assignment.
			_iouMatrix[i][j] = 1 - GetIOU(_predictedBoxes[i], DObj2TBox(retObjs.at(j)));
		}
	}

	// solve the assignment problem using hungarian algorithm.
	// the resulting assignment is [track(prediction) : detection], with len=preNum
	HungarianAlgorithm HungAlgo;
	_assignment.clear();

	if(_iouMatrix.size() > 0 )
		HungAlgo.Solve(_iouMatrix, _assignment);

	// find matches, unmatched_detections and unmatched_predictions
	_unmatchedTrajectories.clear();
	_unmatchedDetections.clear();
	_allItems.clear();
	_matchedItems.clear();

	if (_detNum > _trkNum) //	there are unmatched detections
	{
		for (unsigned int n = 0; n < _detNum; n++)
			_allItems.insert(n);

		for (unsigned int i = 0; i < _trkNum; ++i)
			_matchedItems.insert(_assignment[i]);

		set_difference(_allItems.begin(), _allItems.end(),
			_matchedItems.begin(), _matchedItems.end(),
			insert_iterator<set<int>>(_unmatchedDetections, _unmatchedDetections.begin()));
	}
	else if (_detNum < _trkNum) // there are unmatched trajectory/predictions
	{
		for (unsigned int i = 0; i < _trkNum; ++i)
			if (_assignment[i] == -1) // unassigned label will be set as -1 in the assignment algorithm
				_unmatchedTrajectories.insert(i);
	}
	else
		;

	// filter out matched with low IOU
	_matchedPairs.clear();
	for (unsigned int i = 0; i < _trkNum; ++i)
	{
		if (_assignment[i] == -1) // pass over invalid values
			continue;
		if (1 - _iouMatrix[i][_assignment[i]] < _iouThreshold)
		{
			_unmatchedTrajectories.insert(i);
			_unmatchedDetections.insert(_assignment[i]);
		}
		else
			_matchedPairs.push_back(cv::Point(i, _assignment[i]));
	}
}

void SORTTracker::update_(std::vector<DetectedObject>& retObjs)
{
	///////////////////////////////////////
// 3.3. updating trackers

// update matched trackers with assigned detections.
// each prediction is corresponding to a tracker
	int detIdx, trkIdx;
	for (unsigned int i = 0; i < _matchedPairs.size(); i++)
	{
		trkIdx = _matchedPairs[i].x;
		detIdx = _matchedPairs[i].y;
		_trackers[trkIdx].update(DObj2TBox(retObjs[detIdx]));
	}

	// create and initialise new trackers for unmatched detections
	for (auto umd : _unmatchedDetections)
	{
		KalmanTracker tracker = KalmanTracker(DObj2TBox(retObjs[umd]));
		_trackers.push_back(tracker);
	}

	// get trackers' output
	_frameTrackingResult.clear();
	for (auto it = _trackers.begin(); it != _trackers.end();)
	{
		if (((*it).m_time_since_update < 1) &&
			((*it).m_hit_streak >= _minHits))
		{
			TrackingBox res;
			res.box = (*it).get_state();
			res.id = (*it).m_id + 1;
			_frameTrackingResult.push_back(res);
			it++;
		}
		else
			it++;

		// remove dead tracklet
		if (it != _trackers.end() && (*it).m_time_since_update > _maxAge)
			it = _trackers.erase(it);
	}
}

void SORTTracker::update(std::vector<DetectedObject>& retObjs)
{
	if (_trackers.size() == 0)
	{
		initialize(retObjs);
		predict();
	}
	else
	{
		predict();
		associate(retObjs);
		update_(retObjs);
	}
}

vector<TrackingBox> SORTTracker::getTrackingResults()
{
	return _frameTrackingResult;
}
